pub struct UnionFindSet {
    parent: Vec<usize>,
    rank: Vec<usize>,
}

impl UnionFindSet {
    pub fn new(n: usize) -> Self {
        let parent = (0..n).collect();
        let rank = vec![0; n];
        Self { parent, rank }
    }

    pub fn find(&mut self, x: usize) -> usize {
        if x >= self.parent.len() {
            self.extend(x + 1);
        }
        if self.parent[x] != x {
            self.parent[x] = self.find(self.parent[x]);
        }
        self.parent[x]
    }

    pub fn extend(&mut self, new_size: usize) {
        let n = self.parent.len();
        self.parent.extend(n..new_size);
        self.rank.extend(std::iter::repeat(0).take(new_size));
    }

    pub fn union(&mut self, x: usize, y: usize) {
        let root_x = self.find(x);
        let root_y = self.find(y);
        if root_x != root_y {
            if self.rank[root_x] > self.rank[root_y] {
                self.parent[root_y] = root_x;
            } else {
                self.parent[root_x] = root_y;
                if self.rank[root_x] == self.rank[root_y] {
                    self.rank[root_y] += 1;
                }
            }
        }
    }

    pub fn classes(&mut self) -> usize {
        self.parent
            .iter()
            .enumerate()
            .filter(|(i, &p)| *i == p)
            .count()
    }
}
